#include <iostream>
#include <algorithm>

// Included files for OpenGL Rendering
#ifdef __APPLE__
#include <OpenGL/gl.h>
#else
#include <GL/gl.h>
#endif

#include <cstdlib>
#include <math.h>

#include "photon_mapping.h"
#include "mesh.h"
#include "face.h"
#include "primitive.h"
#include "kdtree.h"
#include "utils.h"
#include "raytracer.h"

// ========================================================================
// DESTRUCTOR
// ========================================================================
PhotonMapping::~PhotonMapping() {
  // cleanup all the photons
  delete kdtree;
}

// ========================================================================
// RECURSIVELY TRACE A SINGLE PHOTON
// ========================================================================
void PhotonMapping::TracePhoton(const Vec3f &position, const Vec3f &direction, 
				const Vec3f &energy, int iter) 
{
  Ray r(position, direction);
  Hit h = Hit();
  Material *m;
  double rf = RandomDouble();
  Vec3f rcolor, dcolor;
  double ravg, davg;
  
  raytracer->CastRay(r, h, false);

  m = h.getMaterial();
  if (!m)
    return;

  /*  Photon p(r.pointAtParameter(h.getT()), -direction, energy, 0);
      kdtree->AddPhoton(p);*/
  
  if (iter > 3)
    return;

  rcolor = m->getReflectiveColor();
  ravg = (rcolor.x()+rcolor.y()+rcolor.z())/3.0;
  dcolor = m->getDiffuseColor();
  davg = (dcolor.x()+dcolor.y()+dcolor.z())/3.0;
  
  if (iter > 0)
    {
      Photon p(r.pointAtParameter(h.getT()), -direction, energy, iter);
      kdtree->AddPhoton(p);
    }


  
  if (rf > 0 && rf <= davg)
    {
      Vec3f color = m->getDiffuseColor();
      Vec3f newdir, newenergy;
      
      newdir = RandomDiffuseDirectionHemisphere(h.getNormal());
      newenergy = energy;
      newenergy.Scale(color.x(),color.y(),color.z());
      newenergy /= davg;
      
      TracePhoton(r.pointAtParameter(h.getT()), newdir, newenergy, iter+1);
    }
  if (rf > davg && rf <= davg+ravg)
    {
      Vec3f newdir, newenergy;
      Vec3f color = m->getReflectiveColor();
      
      newdir = MirrorDirection(h.getNormal(),direction);
      newenergy = energy;
      newenergy.Scale(color.x(),color.y(),color.z());
      newenergy /= ravg;

      TracePhoton(r.pointAtParameter(h.getT()), newdir, newenergy, iter+1);
    }

  return;
}

// ========================================================================
// TRACE THE SPECIFIED NUMBER OF PHOTONS THROUGHOUT THE SCENE
// ========================================================================
void PhotonMapping::TracePhotons() {
  std::cout << "trace photons" << std::endl;

  // first, throw away any existing photons
  delete kdtree;

  // consruct a kdtree to store the photons
  BoundingBox *bb = mesh->getBoundingBox();
  Vec3f min = bb->getMin();
  Vec3f max = bb->getMax();
  Vec3f diff = max-min;
  min -= 0.001*diff;
  max += 0.001*diff;
  kdtree = new KDTree(BoundingBox(min,max));

  // photons emanate from the light sources
  const std::vector<Face*>& lights = mesh->getLights();

  // compute the total area of the lights
  double total_lights_area = 0;
  for (unsigned int i = 0; i < lights.size(); i++) {
    total_lights_area += lights[i]->getArea();
  }

  // shoot a constant number of photons per unit area of light source
  // (alternatively, this could be based on the total energy of each light)
  for (unsigned int i = 0; i < lights.size(); i++) {  
    double my_area = lights[i]->getArea();
    int num = args->num_photons_to_shoot * my_area / total_lights_area;
    // the initial energy for this photon
    Vec3f energy = my_area/double(num) * lights[i]->getMaterial()->getEmittedColor();
    Vec3f normal = lights[i]->computeNormal();
    for (int j = 0; j < num; j++) {
      Vec3f start = lights[i]->RandomPoint();
      // the initial direction for this photon (for diffuse light sources)
      Vec3f direction = RandomDiffuseDirection(normal);
      
      TracePhoton(start,direction,energy,0);
    }
  }
}


// ======================================================================
// PHOTON VISUALIZATION FOR DEBUGGING
// ======================================================================

void PhotonMapping::RenderPhotons() {
  RenderPhotonPositions();
  //  RenderPhotonDirections();
}

// render the position of each photon
void PhotonMapping::RenderPhotonPositions() {  
  if (kdtree == NULL) return;
  glDisable(GL_LIGHTING);
  glPointSize(2);
  glBegin(GL_POINTS);
  std::vector<const KDTree*> todo;  
  todo.push_back(kdtree);
  while (!todo.empty()) {
    const KDTree *node = todo.back();
    todo.pop_back(); 
    if (node->isLeaf()) {
      const std::vector<Photon> &photons = node->getPhotons();
      int num_photons = photons.size();
      for (int i = 0; i < num_photons; i++) {
	const Photon &p = photons[i];
	Vec3f energy = p.getEnergy()*args->num_photons_to_shoot;
	glColor3f(energy.x(),energy.y(),energy.z());
	const Vec3f &position = p.getPosition();
	glVertex3f(position.x(),position.y(),position.z());
      }
    } else {
      todo.push_back(node->getChild1());
      todo.push_back(node->getChild2());
    } 
  }
  glEnd();
  glEnable(GL_LIGHTING);
}


// render the incoming direction of each photon
void PhotonMapping::RenderPhotonDirections() {
  if (kdtree == NULL) return;
  glDisable(GL_LIGHTING);
  glLineWidth(1);
  glBegin(GL_LINES);
  std::vector<const KDTree*> todo;  
  todo.push_back(kdtree);
  BoundingBox *bb = mesh->getBoundingBox();
  double max_dim = bb->maxDim();
  while (!todo.empty()) {
    const KDTree *node = todo.back();
    todo.pop_back(); 
    if (node->isLeaf()) {
      const std::vector<Photon> &photons = node->getPhotons();
      int num_photons = photons.size();
      for (int i = 0; i < num_photons; i++) {
	const Photon &p = photons[i];
	const Vec3f a = p.getPosition();
	Vec3f b = p.getPosition()-(p.getDirectionFrom()*0.02*max_dim);
	Vec3f energy = p.getEnergy()*args->num_photons_to_shoot;
	glColor3f(energy.x(),energy.y(),energy.z());
	glVertex3f(a.x(),a.y(),a.z());
	glVertex3f(b.x(),b.y(),b.z());
      }
    } else {
      todo.push_back(node->getChild1());
      todo.push_back(node->getChild2());
    } 
  }
  glEnd();
  glEnable(GL_LIGHTING);
}


// ======================================================================
// render a simple wireframe of the KD tree
void PhotonMapping::RenderKDTree() {
  if (kdtree == NULL) return;
  glDisable(GL_LIGHTING);
  glLineWidth(1);
  glColor3f(0,0,0);
  glBegin(GL_LINES);
  std::vector<const KDTree*> todo;  
  todo.push_back(kdtree);
  while (!todo.empty()) {
    const KDTree *node = todo.back();
    todo.pop_back(); 
    if (node->isLeaf()) {
      const Vec3f& min = node->getMin();
      const Vec3f& max = node->getMax();
      
      glVertex3f(min.x(),min.y(),min.z());
      glVertex3f(max.x(),min.y(),min.z());
      glVertex3f(min.x(),min.y(),min.z());
      glVertex3f(min.x(),max.y(),min.z());
      glVertex3f(max.x(),max.y(),min.z());
      glVertex3f(max.x(),min.y(),min.z());
      glVertex3f(max.x(),max.y(),min.z());
      glVertex3f(min.x(),max.y(),min.z());

      glVertex3f(min.x(),min.y(),min.z());
      glVertex3f(min.x(),min.y(),max.z());
      glVertex3f(min.x(),max.y(),min.z());
      glVertex3f(min.x(),max.y(),max.z());
      glVertex3f(max.x(),min.y(),min.z());
      glVertex3f(max.x(),min.y(),max.z());
      glVertex3f(max.x(),max.y(),min.z());
      glVertex3f(max.x(),max.y(),max.z());
      
      glVertex3f(min.x(),min.y(),max.z());
      glVertex3f(max.x(),min.y(),max.z());
      glVertex3f(min.x(),min.y(),max.z());
      glVertex3f(min.x(),max.y(),max.z());
      glVertex3f(max.x(),max.y(),max.z());
      glVertex3f(max.x(),min.y(),max.z());
      glVertex3f(max.x(),max.y(),max.z());
      glVertex3f(min.x(),max.y(),max.z());
      
    } else {
      todo.push_back(node->getChild1());
      todo.push_back(node->getChild2());
    } 
  }
  glEnd();
  glEnable(GL_LIGHTING);
}


// ======================================================================

// helper function
bool closest_photon(const std::pair<Photon,double> &a, const std::pair<Photon,double> &b) {
  return (a.second < b.second);
}


// ======================================================================
Vec3f PhotonMapping::GatherIndirect(const Vec3f &point, const Vec3f &normal, const Vec3f &direction_from) const {
  std::vector<Photon> photons;
  if (kdtree == NULL) { 
    std::cout << "WARNING: Photons have not been traced throughout the scene." << std::endl;
    return Vec3f(0,0,0); 
  }
  Vec3f irad(0,0,0);
  int i = 0;
  double bsize = .075;
  double mlength = 0;


  // ================================================================
  // ASSIGNMENT: GATHER THE INDIRECT ILLUMINATION FROM THE PHOTON MAP
  // ================================================================
  
  while (i < args->num_photons_to_collect)
    {
      Vec3f max(point.x()+bsize,point.y()+bsize,point.z()+bsize);
      Vec3f min(point.x()-bsize,point.y()-bsize,point.z()-bsize);
      BoundingBox bbox(min, max);
      
      kdtree->CollectPhotonsInBox(bbox, photons);

      for (std::vector<Photon>::iterator it = photons.begin();
	   it != photons.end();
	   it++)
	{
	  Photon p = *it;

	  if (fabs((p.getPosition()-point).Length()) > bsize)
	    continue;
	  
	  irad += p.getEnergy() * std::max(0.0,normal.Dot3(p.getDirectionFrom()));
	  i++;
	  if (i >= args->num_photons_to_collect)
	    break;
	}

      bsize *= 2;
    }
  irad *= 1.0/(M_PI*bsize*bsize);
  return irad;
}

